from qt import *
from tablewindow import TableWindow
from geticon import get_file_icon
from winappdbg import Process

from pefile import PE
from ctypes import *
from winappdbg.win32.peb_teb import *
from winappdbg.win32.defines import *

import win32process


def get_process_path(pid):
    path = Process(pid).get_filename()
    if not path:
        if pid == 0:
            path = 'System Idle Process'
        elif pid == 4:
            path = 'System'
        else:
            path = '???'
    return path


def may_crash(f):
    def g(*v, **k):
        try:
            return f(*v, **k)
        except Exception, e:
            print e
    return g

    
def suggest_dumped_name(path):
    import os
    head, tail = os.path.split(path)
    return os.path.join(head, '_' + tail)

class ProcIndex:
    Path, Pid, Base, Size = range(4)
    
class ModIndex:
    Path, Base, Size = range(3)


def fix_dumped(data):
    """ Set FileOffset=RVA for each section, return as a string """
    pe = PE(data=data)
    for section in pe.sections:
        section.PointerToRawData = section.VirtualAddress
        section.Misc_VirtualSize = section.SizeOfRawData
    return pe.write()
    

def dump_module(module):
    process = module.get_process()
    defpath = suggest_dumped_name(module.get_filename())
    path = QFileDialog.getSaveFileName(directory=defpath)
    if path:
        base = module.get_base()
        size = module.get_size()
        data = process.read(base, size)
        data = fix_dumped(data)
        with open(path, 'wb') as ofile:
            ofile.write(data)
                
                    

class ProcessWindow(QMainWindow):

    def __init__(self, parent=None):
        super(ProcessWindow, self).__init__(parent)
        self.setWindowTitle('Processes')
        spl = QSplitter(self)

        self.process_list_area = TableWindow(self, ['Image Path', 'PID', 'ImageBase', 'SizeOfImage'])
        self.process_list_area.table.setColumnWidth(0, 512)
        self.process_list_area.table.setContextMenuPolicy(Qt.CustomContextMenu)
        self.process_list_area.table.customContextMenuRequested.connect(self.process_menu)
        self.process_list_area.table.itemSelectionChanged.connect(self.selected_process_changed)
        
        self.module_list_area = TableWindow(self, ['Module Path', 'ImageBase', 'SizeOfImage'])
        self.module_list_area.table.setColumnWidth(0, 512)
        self.module_list_area.table.setContextMenuPolicy(Qt.CustomContextMenu)
        self.module_list_area.table.customContextMenuRequested.connect(self.module_menu)
        self.module_list_area.table.itemSelectionChanged.connect(self.selected_module_changed)
                
        spl.addWidget(self.process_list_area)
        spl.addWidget(self.module_list_area)
        spl.setOrientation(Qt.Vertical)
        self.setCentralWidget(spl)        

        self.flush_process_table()
        
    def process_menu(self, point):
        menu = QMenu(self)
        menu.addAction(QAction('dump full', self, triggered=self.dump_process_full))
        menu.addAction(QAction('dump region', self, triggered=self.dump_process_region, enabled=False))
        menu.addSeparator()
        menu.addAction(QAction('inject dll', self, triggered=self.inject_dll))
        menu.addAction(QAction('correct SizeOfImage', self, triggered=self.correct_SizeOfImage))
        menu.addSeparator()
        menu.addAction(QAction('burn process', self, triggered=self.burn_process))
        menu.addSeparator()
        menu.addAction(QAction('refresh', self, triggered=self.flush_process_table))
        menu.popup(QCursor.pos())
        
    def module_menu(self, point):
        menu = QMenu(self)
        menu.addAction(QAction('dump module', self, triggered=self.dump_module))
        menu.addSeparator()
        menu.addAction(QAction('correct SizeOfImage', self, triggered=self.correct_SizeOfImage_module))
        menu.addAction(QAction('unlink in Ldr', self, triggered=self.unlink_ldr, enabled=False))
        menu.popup(QCursor.pos())
        

    def flush_process_table(self):
        self.process_list_area.delete_all_items()
    
        for pid in reversed(win32process.EnumProcesses()):
            process = Process(pid)
            
            try:
                module = process.get_main_module()
                base = module.get_base()
                size = module.get_size()
            except:
                base = size = 0
            path = get_process_path(pid)
            item = QTableWidgetItem(path)
            icon = get_file_icon(path)
            if icon:
                item.setIcon(icon)
            self.process_list_area.add_item(
                item, 
                '%08X' % pid,
                '%08X' % (base),
                '%08X' % (size),
            )

    @may_crash
    def flush_module_table(self, pid):
        self.module_list_area.delete_all_items()

        process = Process(pid)
        for module in process.iter_modules():
            path = module.get_filename()
            item = QTableWidgetItem(path)
            icon = get_file_icon(path)
            if icon:
                item.setIcon(icon)
            self.module_list_area.add_item(
                item,
                '%08X' % module.get_base(),
                '%08X' % module.get_size(),
            )

    def get_selected_pid(self):
        items = self.process_list_area.table.selectedItems()
        if not items:
            return
        pid = items[ProcIndex.Pid]
        pid = str(pid.text())
        pid = int(pid, 16)
        return pid

    def get_selected_process(self):
        return Process(self.get_selected_pid())
        
    def get_selected_path(self):
        items = self.process_list_area.table.selectedItems()
        if not items:
            return
        path = items[ProcIndex.Path]
        return unicode(path.text())

    def get_selected_module(self):
        process = self.get_selected_process()
        items = self.module_list_area.table.selectedItems()        
        base = items[ModIndex.Base]
        base = str(base.text())
        base = int(base, 16)        
        module = process.get_module(base)
        return module
        
    def selected_process_changed(self):
        pid = self.get_selected_pid()
        self.flush_module_table(pid)

    def selected_module_changed(self):
        pass

    def burn_process(self):
        process = self.get_selected_process()
        process.kill()
        
    def dump_process_full(self):
        process = self.get_selected_process()
        module = process.get_main_module()
        dump_module(module)
        
    def dump_process_region(self):
        1/0

    def dump_module(self):
        module = self.get_selected_module()
        dump_module(module)

    def inject_dll(self):
        dll = QFileDialog.getOpenFileName()
        process = self.get_selected_process()
        if process and dll:
            process.inject_dll(dll)
            
        
    def correct_SizeOfImage(self):
        process = self.get_selected_process()
        module = process.get_main_module()
        fix_image_size(module)
        self.flush_process_table()

    def correct_SizeOfImage_module(self):
        fix_image_size(self.get_selected_module())
        self.flush_module_table()
        
    def unlink_ldr(self):
        1/0

def CONTAINING_RECORD(address, type, field):
    return address - getattr(type, field).offset
    
    
def read_linked_list(process, header, type, field):
    address = header
    while True:
        record = type()
        data = process.read(CONTAINING_RECORD(address, type, field), sizeof(record))
        memmove(byref(record), data, sizeof(record))
        
        yield address, record
        
        address = getattr(record, field).Flink
        if address == header:
            break

def get_ldr(process):
    peb = process.get_peb()
    ldr = PEB_LDR_DATA()
    memmove(byref(ldr), process.read(peb.Ldr, sizeof(ldr)), sizeof(ldr))
    return ldr
    
def fix_image_size(module):
    process = module.get_process()
    filename = module.get_filename()
    pe = PE(filename)
    realsize = pe.OPTIONAL_HEADER.SizeOfImage
    
    ldr = get_ldr(process)
    def handle_ldr_module(address, record):
        if record.BaseAddress == module.get_base():
            if record.SizeOfImage != realsize:
                print ('fix image size for module %s %08X -> %08X' % (module.get_filename(), record.SizeOfImage, realsize))
                record.SizeOfImage = realsize
                process.write(address, buffer(record)[:sizeof(record)])
                
    for address, record in read_linked_list(process, ldr.InInitializationOrderModuleList.Flink, LDR_MODULE, 'InInitializationOrderModuleList'):
        handle_ldr_module(address, record)
    for address, record in read_linked_list(process, ldr.InLoadOrderModuleList.Flink, LDR_MODULE, 'InLoadOrderModuleList'):
        handle_ldr_module(address, record)
    for address, record in read_linked_list(process, ldr.InMemoryOrderModuleList .Flink, LDR_MODULE, 'InMemoryOrderModuleList'):
        handle_ldr_module(address, record)        
    
    
if __name__ == '__main__':
    from winappdbg import System
    System.request_debug_privileges()
    
    app = QApplication([])
    wnd = ProcessWindow()
    wnd.show()
    app.exec_()